//! 地址转换
//!
//! # TODO
//! 当前设计中无法处理内存空洞问题, 因此假定只有一块内存域来实现下面的功能

use core::{
    mem::size_of,
    ops::{Add, AddAssign, BitAnd, BitAndAssign, Sub},
};

use super::kalloc::Page;
use crate::{
    bsp::__mems_define,
    println,
    space::{
        mm::pgtabledef::{PAGE_SHIFT, PAGE_SIZE},
        uaccess::USER_SPACE_SIZE,
    },
};

cfg_if::cfg_if! {
    if #[cfg(target_arch = "aarch64")] {
        mod aarch64;
        use aarch64::*;
    } else {
        mod dummy;
        use dummy::*;
    }
}

// init task 栈基址
#[inline]
pub(crate) fn init_stack_base() -> Kaddr {
    Kaddr(arch_init_stack_base())
}

// 内核代码的起始偏移地址, 可实现物理地址和内核地址的互相转换
const KIMAGE_OFFSET: usize = ARCH_KIMAGE_OFFSET;
// 线性内存起始偏移地址, 可实现物理地址和线性地址的互相转换
const KMEM_OFFSET: usize = ARCH_KMEM_OFFSET;
// 页结构地址偏移, 可用于通过其他地址快速找到对应的页结构
const KPAGE_OFFSET: usize = ARCH_KPAGE_OFFSET;
// 内核 io 起始偏移地址, 可实现物理地址和 io 地址的互相转换
const KIO_OFFSET: usize = ARCH_KIO_OFFSET;

// 内核起始虚拟镜像地址
#[inline]
fn start_text() -> Kaddr {
    Kaddr(arch_start_text())
}

/// 内核结束虚拟镜像地址
#[inline]
fn end_bss() -> Kaddr {
    Kaddr(arch_end_bss() - 1)
}

// 内存起始物理地址
#[inline(always)]
pub(crate) fn memstart() -> Paddr {
    unsafe { Paddr(__mems_define().mem(0).0) }
}

// 内存大小
#[inline(always)]
fn memsize() -> usize {
    debug_assert!(unsafe { __mems_define().nr_mems() == 1 });
    unsafe { __mems_define().mem(0).1 }
}

// 内存结束物理地址
#[inline(always)]
fn memend() -> Paddr {
    memstart() + memsize() - 1
}

// 页结构大小
const KPAGE_SIZE: usize = size_of::<Page>();
// 页结构总数量
#[inline(always)]
fn kpage_count() -> usize {
    memsize() / PAGE_SIZE
}
// 页结构占用总内存量
#[inline(always)]
fn kpage_mem() -> usize {
    KPAGE_SIZE * kpage_count()
}

/// 页结构起始物理地址
pub fn kpage_start_phys() -> usize {
    memend().to_value() + 1 - kpage_mem()
}

// 实际的页结构起始地址
// 通过定义为物理地址高位位置到内存结束区域为页结构存储空间
#[inline(always)]
fn kpage_start() -> usize {
    KPAGE_OFFSET + kpage_start_phys()
}

// 页结构结束地址
#[inline(always)]
fn kpage_end() -> usize {
    kpage_start() + kpage_mem() - 1
}

/// 物理地址结构
#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
pub struct Paddr(usize);

impl Add<usize> for Paddr {
    type Output = Paddr;

    #[inline(always)]
    fn add(self, rhs: usize) -> Self::Output {
        Paddr(self.0 + rhs)
    }
}

impl Sub<usize> for Paddr {
    type Output = Paddr;

    #[inline(always)]
    fn sub(self, rhs: usize) -> Self::Output {
        Paddr(self.0 - rhs)
    }
}

impl Sub for Paddr {
    type Output = usize;

    #[inline(always)]
    fn sub(self, rhs: Self) -> Self::Output {
        self.0 - rhs.0
    }
}

impl Paddr {
    /// 常量构造
    pub const fn from(value: usize) -> Paddr {
        Paddr(value)
    }

    /// 物理地址值
    pub const fn to_value(&self) -> usize {
        self.0
    }

    /// 物理地址转换为内核地址
    #[inline(always)]
    pub fn to_kimg(&self) -> Option<Kaddr> {
        let value = self.0 + KIMAGE_OFFSET;
        if value >= start_text().to_value() && value <= end_bss().to_value() {
            Some(Kaddr(value))
        } else {
            None
        }
    }

    /// 物理地址转换为线性地址
    #[inline(always)]
    pub fn to_virt(&self) -> Option<Vaddr> {
        if self.phys_valid() {
            return Some(Vaddr(self.0 + KMEM_OFFSET));
        }
        None
    }

    /// 物理地址转换为页
    ///
    /// # TODO
    /// 这里的设计需要保证内存是连续的, 没有空洞内存,
    /// 因为 KPAGE 是连续内存映射的.
    #[inline(always)]
    pub fn to_page(&self) -> Option<&'static Page> {
        let idx = self.to_pfn()?;
        unsafe {
            let page = (kpage_start() as *const Page).add(idx.to_value());
            if page as usize > kpage_end() { None } else { Some(&*page) }
        }
    }

    /// 物理地址转换为可变页
    #[inline(always)]
    pub fn to_page_mut(&self) -> Option<&'static mut Page> {
        let idx = self.to_pfn()?;
        unsafe {
            let page = (kpage_start() as *mut Page).add(idx.to_value());
            if page as usize > kpage_end() { None } else { Some(&mut *page) }
        }
    }

    /// 物理地址转换为IO地址
    #[inline(always)]
    pub fn to_io(&self) -> Option<Iaddr> {
        if self.0 >= memstart().to_value() && self.0 <= memend().to_value() {
            None
        } else {
            Some(Iaddr(self.0 + KIO_OFFSET))
        }
    }

    /// 物理地址常量转换不做验证
    ///
    /// # NOTES
    /// 仅用于BSP初始化
    pub const fn to_io_const(&self) -> Iaddr {
        Iaddr(self.0 + KIO_OFFSET)
    }

    /// 物理地址转换为页帧
    #[inline(always)]
    pub fn to_pfn(&self) -> Option<Pfn> {
        if self.0 < memstart().to_value() || self.0 > memend().to_value() {
            None
        } else {
            Some(Pfn((self.0 - memstart().to_value()) >> PAGE_SHIFT))
        }
    }

    /// 物理地址向上对齐转换为页帧
    #[inline(always)]
    pub fn to_pfn_up(&self) -> Option<Pfn> {
        let value = self.0 + PAGE_SIZE - 1;
        Paddr::from(value).to_pfn()
    }

    /// 物理地址是否是空闲可用内存区域
    pub fn phys_valid(&self) -> bool {
        unsafe {
            let index = MEMBLOCK_RANGE.index;
            for i in 0..index {
                let start = MEMBLOCK_RANGE.range[i].start;
                let end = start + MEMBLOCK_RANGE.range[i].size;
                if *self >= start && *self < end {
                    return true;
                }
            }
        }
        false
    }
}

/// 内核地址结构
pub struct Kaddr(usize);

impl From<usize> for Kaddr {
    #[inline(always)]
    fn from(value: usize) -> Self {
        assert!(value >= start_text().to_value() && value <= end_bss().to_value());
        Kaddr(value)
    }
}

impl Kaddr {
    /// 内核地址值
    pub const fn to_value(&self) -> usize {
        self.0
    }

    /// 内核地址转换为物理地址
    #[inline(always)]
    pub fn to_phys(&self) -> Paddr {
        debug_assert!(self.0 >= start_text().to_value() && self.0 <= end_bss().to_value());
        Paddr(self.0 - KIMAGE_OFFSET)
    }

    /// 内核地址转换为页
    ///
    /// # Panics
    /// 系统故障将会 panic.
    #[inline(always)]
    pub fn to_page(&self) -> &'static Page {
        self.to_phys().to_page().unwrap()
    }
}

/// 线性地址结构
#[derive(Clone, Copy)]
pub struct Vaddr(usize);

impl From<usize> for Vaddr {
    #[inline(always)]
    fn from(value: usize) -> Self {
        let vaddr = Vaddr(value);
        debug_assert!(vaddr.to_phys().phys_valid());
        vaddr
    }
}

impl Vaddr {
    /// 线性地址值
    pub const fn to_value(&self) -> usize {
        self.0
    }

    /// 线性地址转换为物理地址
    ///
    /// # Panics
    /// 系统故障将会 panic.
    #[inline(always)]
    pub fn to_phys(&self) -> Paddr {
        let phys = Paddr(self.0 - KMEM_OFFSET);
        debug_assert!(phys.phys_valid());
        phys
    }

    /// 线性地址转换为页
    ///
    /// # Panics
    /// 系统故障将会 panic.
    #[inline(always)]
    pub fn to_page(&self) -> &'static Page {
        self.to_phys().to_page().unwrap()
    }

    /// 线性地址转换为可变页
    ///
    /// # Panics
    /// 系统故障将会 panic.
    #[inline(always)]
    pub fn to_page_mut(&self) -> &'static mut Page {
        self.to_phys().to_page_mut().unwrap()
    }

    /// 线性地址转换为页帧
    ///
    /// # Panics
    /// 系统故障将会 panic.
    #[inline(always)]
    pub fn to_pfn(&self) -> Pfn {
        self.to_phys().to_pfn().unwrap()
    }
}

impl Page {
    /// 页转换为物理地址
    #[inline(always)]
    pub fn to_phys(&self) -> Paddr {
        let p = self as *const Page as usize;
        debug_assert!(p >= kpage_start() && p <= kpage_end());
        Paddr((((p - kpage_start()) / KPAGE_SIZE) << PAGE_SHIFT) + memstart().to_value())
    }

    /// 页转换为内核地址
    #[inline(always)]
    pub fn to_kimg(&self) -> Option<Kaddr> {
        self.to_phys().to_kimg()
    }

    /// 页转换为线性地址
    #[inline(always)]
    pub fn to_virt(&self) -> Option<Vaddr> {
        self.to_phys().to_virt()
    }

    /// 页转换页帧
    ///
    /// # Panics
    /// 系统故障将会 panic.
    #[inline(always)]
    pub fn to_pfn(&self) -> Pfn {
        self.to_phys().to_pfn().unwrap()
    }
}

/// IO地址结构
pub struct Iaddr(usize);

impl Iaddr {
    /// IO地址值
    #[inline(always)]
    pub const fn to_value(&self) -> usize {
        self.0
    }

    /// IO地址转换为物理地址
    #[inline(always)]
    pub fn to_phys(&self) -> Paddr {
        let phys = Paddr(self.0 - KIO_OFFSET);
        debug_assert!(phys < memstart() || phys > memend());
        phys
    }
}

/// 页帧结构
#[derive(Clone, Copy, PartialEq, PartialOrd)]
pub struct Pfn(usize);

impl From<usize> for Pfn {
    #[inline(always)]
    fn from(value: usize) -> Self {
        debug_assert!((memstart() + (value << PAGE_SHIFT)) <= memend());
        Pfn(value)
    }
}

impl Add<usize> for Pfn {
    type Output = Pfn;

    #[inline(always)]
    fn add(self, rhs: usize) -> Self::Output {
        Self(self.0 + rhs)
    }
}

impl Sub for Pfn {
    type Output = usize;

    #[inline(always)]
    fn sub(self, rhs: Self) -> Self::Output {
        self.0 - rhs.0
    }
}

impl Sub<usize> for Pfn {
    type Output = Pfn;

    #[inline(always)]
    fn sub(self, rhs: usize) -> Self::Output {
        Pfn(self.0 - rhs)
    }
}

impl BitAnd for Pfn {
    type Output = Pfn;

    #[inline(always)]
    fn bitand(self, rhs: Self) -> Self::Output {
        Pfn(self.0 & rhs.0)
    }
}

impl Pfn {
    // 构造一个错误的 pfn
    pub(crate) const fn from_error() -> Self {
        Pfn(usize::MAX)
    }

    /// 页帧值
    pub const fn to_value(&self) -> usize {
        self.0
    }

    /// 页帧转换为物理地址
    #[inline(always)]
    pub fn to_phys(&self) -> Paddr {
        let phys = memstart() + (self.0 << PAGE_SHIFT);
        debug_assert!(phys <= memend());
        phys
    }

    /// 页帧转换为页
    #[inline(always)]
    pub fn to_page(&self) -> Option<&'static Page> {
        self.to_phys().to_page()
    }

    /// 页帧转换为可变页
    #[inline(always)]
    pub fn to_page_mut(&self) -> Option<&'static mut Page> {
        self.to_phys().to_page_mut()
    }

    /// 页帧是否在线性区域
    #[inline(always)]
    pub fn pfn_valid(&self) -> bool {
        let phys = self.to_phys();
        phys.phys_valid()
    }
}

/// 应用虚拟地址
#[derive(Clone, Copy, PartialEq, PartialOrd, Eq, Ord)]
pub struct Uaddr(usize);

impl From<usize> for Uaddr {
    #[inline(always)]
    fn from(value: usize) -> Self {
        debug_assert!(value <= USER_SPACE_SIZE);
        Uaddr(value)
    }
}

impl Sub<usize> for Uaddr {
    type Output = Uaddr;

    #[inline(always)]
    fn sub(self, rhs: usize) -> Self::Output {
        Uaddr(self.0 - rhs)
    }
}

impl Sub for Uaddr {
    type Output = usize;

    #[inline(always)]
    fn sub(self, rhs: Self) -> Self::Output {
        self.0 - rhs.0
    }
}

impl Add<usize> for Uaddr {
    type Output = Uaddr;

    #[inline(always)]
    fn add(self, rhs: usize) -> Self::Output {
        Uaddr(self.0 + rhs)
    }
}

impl BitAnd<usize> for Uaddr {
    type Output = Uaddr;

    #[inline(always)]
    fn bitand(self, rhs: usize) -> Self::Output {
        Uaddr(self.0 & rhs)
    }
}

impl BitAndAssign<usize> for Uaddr {
    #[inline(always)]
    fn bitand_assign(&mut self, rhs: usize) {
        self.0 = self.0 & rhs;
    }
}

impl AddAssign<usize> for Uaddr {
    #[inline(always)]
    fn add_assign(&mut self, rhs: usize) {
        self.0 += rhs;
    }
}

impl Uaddr {
    /// 应用虚拟地址值
    pub const fn to_value(&self) -> usize {
        self.0
    }
}

#[derive(Clone, Copy)]
struct MemBlockRange {
    start: Paddr,
    size: usize,
}

// 定义有 16 块, 目前足够
const MEM_RANG_MAX: usize = 16;

struct MemBlock {
    index: usize,
    range: [MemBlockRange; MEM_RANG_MAX],
}

// 定义一个全局的可用内存范围, 最多可以设置 MEM_RANG_MAX 块内存区域
static mut MEMBLOCK_RANGE: MemBlock =
    MemBlock { index: 0, range: [MemBlockRange { start: Paddr(0), size: 0 }; MEM_RANG_MAX] };

// 添加一块内存域
//
// 注意: 添加内存域不能与已添加的内存域交叉并且不能添加超过 `MEM_RANG_MAX`, 否则将会断言失败
fn memblock_add(start: Paddr, size: usize) {
    let end = start + size;
    unsafe {
        let index = MEMBLOCK_RANGE.index;
        assert!(index < MEM_RANG_MAX);
        for i in 0..index {
            let mem_end = MEMBLOCK_RANGE.range[i].start + MEMBLOCK_RANGE.range[i].size;
            assert!(end <= MEMBLOCK_RANGE.range[i].start || mem_end <= start);
        }
        MEMBLOCK_RANGE.range[index].start = start;
        MEMBLOCK_RANGE.range[index].size = size;
        MEMBLOCK_RANGE.index += 1;
    }
}

// 遍历内存域
pub(crate) fn memblock_foreach<F>(f: F)
where
    F: Fn(Paddr, usize, usize),
{
    unsafe {
        let index = MEMBLOCK_RANGE.index;
        for i in 0..index {
            let start = MEMBLOCK_RANGE.range[i].start;
            let size = MEMBLOCK_RANGE.range[i].size;
            f(start, size, i);
        }
    }
}

// 全局可用内存域设置
//
// f: start, size, `is_free`
pub(crate) fn mem_setup<F>(f: F)
where
    F: Fn(Paddr, usize, bool),
{
    // 向下对齐结束地址
    let end = start_text().to_phys().to_pfn().unwrap().to_phys();
    let size = end - memstart();
    memblock_add(memstart(), size);
    f(memstart(), size, true);

    let start = end;
    let end = end_bss().to_phys().to_pfn_up().unwrap().to_phys();
    f(start, end - start, false);

    let start = end;
    let end = Paddr(kpage_start_phys()).to_pfn().unwrap().to_phys();
    memblock_add(start, end - start);
    f(start, end - start, true);
    f(end, memend() + 1 - end, false);
}

/// dump 内存分布
///
/// # Panics
/// 系统故障将会 panic.
pub fn mem_dump() {
    println!("phys:  [{:#018x}-{:#018x}]", memstart().to_value(), memend().to_value());
    println!("kimg:  [{:#018x}-{:#018x}]", start_text().to_value(), end_bss().to_value());
    println!("page:  [{:#018x}-{:#018x}]", kpage_start(), kpage_end());
    memblock_foreach(|start, size, idx| {
        println!(
            "virt{}: [{:#018x}-{:#018x}]",
            idx,
            start.to_virt().unwrap().to_value(),
            (start + size - 1).to_virt().unwrap().to_value()
        );
    });
}
